#include <iostream>
#include <fstream>
#include <vector>
#include <random>
#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm-c/Core.h"
#include "llvm/ADT/APInt.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
using namespace llvm;

llvm::LLVMContext *CONTEXT=nullptr;


namespace llvm
{
    BasicBlock* createCloneBasicBlock(BasicBlock *BB);

    struct RandomControlFlow : public FunctionPass
    {
    
        static char ID;
        RandomControlFlow() : FunctionPass(ID) {}
        bool runOnFunction(Function &F);
        Value* opaquePredicates(Value *x,BasicBlock *BB);

        void insertFakeFlow(Value *x, BasicBlock *ifTrue, BasicBlock *ifFalse, BasicBlock *BB);

        bool randControlFlow(BasicBlock *BB);
    };
}
char RandomControlFlow::ID = 0;
static RegisterPass<RandomControlFlow> X("randCF", "insert fake random blocks");
Pass *createRandomControlFlowPass() { return new RandomControlFlow(); }

//ollvm
BasicBlock* llvm::createCloneBasicBlock(BasicBlock *BB)
{
    std::vector<Instruction*> origReg;
    BasicBlock &entryBB = BB->getParent()->getEntryBlock();
    for(Instruction &I : *BB)
    {
        if(!(isa<AllocaInst>(&I) && I.getParent() == &entryBB) && I.isUsedOutsideOfBlock(BB))
        {
            origReg.push_back(&I);
        }
    }
    for(Instruction *I : origReg)
    {
        DemoteRegToStack(*I, entryBB.getTerminator());
    }
    ValueToValueMapTy VMap;
    BasicBlock *cloneBB = CloneBasicBlock(BB, VMap, "cloneBB", BB->getParent());
    for(Instruction &I : *cloneBB)
    {
        for(int i=0;i<I.getNumOperands(); i++)
        {
            Value *V = MapValue(I.getOperand(i), VMap);
            if(V)
            {
                I.setOperand(i, V);
            }
        }
    }
    return cloneBB;
}

bool RandomControlFlow::runOnFunction(Function &F)
{
    CONTEXT = &F.getContext();
    for (int i=0;i<2;i++)
    {
        std::vector<BasicBlock*> origBB;
        for(BasicBlock &BB : F)
        {
            origBB.push_back(&BB);
        }
        for(BasicBlock *BB : origBB)
        {
            randControlFlow(BB);
        }
    }
    return true;
}

void RandomControlFlow::insertFakeFlow(Value *x,
                    BasicBlock *ifTrue, BasicBlock *ifFalse, BasicBlock *BB)
{
    Value *fakeExpr = opaquePredicates(x, BB);
    //Value *randMod2 = BinaryOperator::Create(Instruction::And,
    //                                alteredRandVar, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 1, false), "", insertAfter);
    ICmpInst *condition = new ICmpInst(*BB,
                                    ICmpInst::ICMP_EQ, fakeExpr, x);
    BranchInst::Create(ifTrue, ifFalse, condition, BB);
}

bool RandomControlFlow::randControlFlow(BasicBlock *BB)
{

    BasicBlock *entryBB = BB;
    BasicBlock *bodyBB = entryBB->splitBasicBlock(BB->getFirstNonPHIOrDbgOrLifetime(), "bodyBB");
    BasicBlock *endBB = bodyBB->splitBasicBlock(bodyBB->getTerminator(), "endBB");
    BasicBlock *cloneBB = createCloneBasicBlock(bodyBB);
    
    entryBB->getTerminator()->eraseFromParent();

    Module &M = *BB->getModule();
    Function *F =
    cast<Function>(M.getOrInsertFunction("rand",
                                FunctionType::getInt32Ty(*CONTEXT)).getCallee());
    F->setDSOLocal(true);
    Value *x = CallInst::Create(F->getFunctionType(),
                                      F, "", entryBB);

    insertFakeFlow(x, bodyBB, cloneBB, entryBB);

    bodyBB->getTerminator()->eraseFromParent();
    insertFakeFlow(x, endBB, cloneBB, bodyBB);
    cloneBB->getTerminator()->eraseFromParent();
    insertFakeFlow(x, bodyBB, endBB, cloneBB);

    return true;
}

Value* RandomControlFlow::opaquePredicates(Value *x,BasicBlock *BB)
{
    uint32_t tmp = rand() % 3;
    Value *ret;
    if(tmp == 0)
    {

        //x*(x^2-x+1)-x^3+x^2
        BinaryOperator *op1 = BinaryOperator::Create(Instruction::Mul, x, x, "", BB);
        BinaryOperator *op2 = BinaryOperator::Create(Instruction::Sub, op1, x, "", BB);
        BinaryOperator *op3 = BinaryOperator::Create(Instruction::Add, op2, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 1, false), "", BB);
        BinaryOperator *op4 = BinaryOperator::Create(Instruction::Mul, op3, x, "", BB);
        
        BinaryOperator *op5 = BinaryOperator::Create(Instruction::Mul, x, x, "", BB);
        BinaryOperator *op6 = BinaryOperator::Create(Instruction::Mul, op5, x, "", BB);
        BinaryOperator *op7 = BinaryOperator::Create(Instruction::Sub, op4, op6, "", BB);
        BinaryOperator *op8 = BinaryOperator::Create(Instruction::Add, op7, op5, "", BB);
        ret = op8;
    }else
        if(tmp == 1)
        {
            //2x *(x + 3) - 2x^2 - 5 * x
            BinaryOperator *op1 = BinaryOperator::Create(Instruction::Mul, x, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 2, false), "", BB);
            BinaryOperator *op2 = BinaryOperator::Create(Instruction::Add, x, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 3, false), "", BB);
            BinaryOperator *op3 = BinaryOperator::Create(Instruction::Mul, op1, op2, "", BB);
            
            BinaryOperator *op4 = BinaryOperator::Create(Instruction::Mul, x, x, "", BB);
            BinaryOperator *op5 = BinaryOperator::Create(Instruction::Mul, op4, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 2, false), "", BB);
            
            BinaryOperator *op6 = BinaryOperator::Create(Instruction::Mul, x, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 5, false), "", BB);
            
            BinaryOperator *op7 = BinaryOperator::Create(Instruction::Sub, op3, op5, "", BB);
            BinaryOperator *op8 = BinaryOperator::Create(Instruction::Sub, op6, op7, "", BB);
            ret = op8;
    }else
        if(tmp == 2)
        {
            //x = (x+1)x-x^2
            BinaryOperator *op1 = BinaryOperator::Create(Instruction::Add, x, ConstantInt::get(Type::getInt32Ty(*CONTEXT), 1, false), "", BB);
            BinaryOperator *op2 = BinaryOperator::Create(Instruction::Mul, op1, x, "", BB);
            BinaryOperator *op3 = BinaryOperator::Create(Instruction::Mul, x, x, "", BB);
            
            BinaryOperator *op4 = BinaryOperator::Create(Instruction::Sub, op2, op3, "", BB);
            ret = op4;
    }
    return ret;
}



